/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package pl.mwas.proteana.counters;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.HashMap;
import pl.mwas.proteana.annotations.AnnotationDb;
import pl.mwas.proteana.annotations.GOAnnotation;
import pl.mwas.proteana.file.readers.ProteanaFileReader;
import pl.mwas.proteana.geneontology.GOTerm;
import pl.mwas.proteana.geneontology.GOTermDb;
import pl.mwas.proteana.geneontology.util.Relation;

/**
 *
 * @author mateusz.was
 */
public class TermCounter {

    private ArrayList<String> input;
    private HashMap<GOTerm, BigDecimal> result;
    private AnnotationDb annotations;
    private GOTermDb ontology;

    private boolean annotationsOnly;

    public TermCounter() {
        result = new HashMap<>();
        input = new ArrayList<>();
    }

    public TermCounter(String path) {
        this();
        input = new ProteanaFileReader(path).getOutput();
    }

    public void setAnnotations(AnnotationDb annotations) {
        this.annotations = annotations;
    }

    public void setOntology(GOTermDb ontology) {
        this.ontology = ontology;
    }

    public void setAnnotationsOnly(boolean annotationsOnly) {
        this.annotationsOnly = annotationsOnly;
    }

    private void countAnnotationsOnly() {
        for (String id : input) {
            ArrayList<GOAnnotation> currAnnotationList = annotations.getAnnotationsForId(id);
            if (currAnnotationList == null) {
                System.err.println("No annotation for id: " + id);
            } else {
                for (GOAnnotation annot : currAnnotationList) {
                    addGoTermToResult(annot.getGoTerm());
                }
            }
        }
    }

    private void addGoTermToResult(GOTerm currGoTerm) {
        if (result.containsKey(currGoTerm)) {
            BigDecimal currCount = result.get(currGoTerm);
            result.put(currGoTerm, currCount.add(BigDecimal.ONE));
        } else {
            result.put(currGoTerm, BigDecimal.ONE);
        }
    }

    public void linkAnnotationWithGoTerms() {
        annotations.addGoTermsToAnnotations(ontology);
    }

    public HashMap<GOTerm, BigDecimal> countTerms() {
        if(annotationsOnly) {
            countAnnotationsOnly();
        }
        else {
            countIncludingRelations();
        }

        return result;
    }

    private void countIncludingRelations() {
        for (String id : input) {
            ArrayList<GOAnnotation> currAnnotationList = annotations.getAnnotationsForId(id);
            if (currAnnotationList == null) {
                System.err.println("No annotation for id: " + id);
            } else {
                for (GOAnnotation annot : currAnnotationList) {
                    GOTerm currGOTerm = annot.getGoTerm();
                    countRelations(currGOTerm);
                }
            }
        }
    }

    private void countRelations(GOTerm currGOTerm) {
        ArrayList<Relation> relationships = currGOTerm.getRelationships();
        if(relationships != null) {
            for(Relation rel : relationships) {
                //TODO not all relation types needs to be count
                countRelations(rel.getRelatedTerm());
                addGoTermToResult(rel.getRelatedTerm());
            }
        }
        else {
            addGoTermToResult(currGOTerm);
        }
    }
}
